import torch
import numpy as np
import argparse
import os
import cv2
from tqdm import tqdm
import random

from data.utils import read_image
from data.dataset import build_dataset
from engine import Predictor
from utils.config import get_test_cfg
from utils.boxes import postprocess


def parse_args():
    parser = argparse.ArgumentParser(description="demo")
    parser.add_argument("--test-datasets", type=str, help='datasets for sampling', required=True)
    parser.add_argument('-n', "--num", type=int, help='number of images', required=True)
    args = parser.parse_args()
    return args

class TinyDataset:
    def __init__(self, dataset_names, max_size):
        self.dataset_dicts = build_dataset(*dataset_names)
        self.dataset_dicts = random.sample(self.dataset_dicts, max_size)

    def __len__(self):
        return len(self.dataset_dicts)

    def __getitem__(self, index):
        filename = self.dataset_dicts[index]['file_name']
        image = read_image(filename, preprocess_method=self.dataset_dicts[index]['preprocess'], params=self.dataset_dicts[index]['params'])
        gt_boxes = [obj['bbox'] for obj in self.dataset_dicts[index]['annotations']]
        return filename, image, gt_boxes
    
    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

def draw_image(img, gt_boxes):
    num_boxes = len(gt_boxes)
    for i in range(num_boxes):
        cv2.rectangle(img, (int(gt_boxes[i][0]), int(gt_boxes[i][1])), (int(gt_boxes[i][2]), int(gt_boxes[i][3])), (0, 255, 0), 3)

    return img
    

if __name__ == '__main__':
    args = parse_args()
    
    test_datasets = ["{}-test".format(dataset) for dataset in args.test_datasets.split(',')]
    
    dataset = TinyDataset(test_datasets, args.num)
    
    output_dir = os.path.join('output','draw_dataset', '{}'.format('&'.join(test_datasets)))
    os.makedirs(output_dir, exist_ok=True)

    count = 0
    for filename, image, gt_boxes in tqdm(dataset):
        count += 1
        image = cv2.imread(filename)
        image = draw_image(image, gt_boxes)
        output_path = os.path.join(output_dir, '{:0>8}{}'.format(count, os.path.splitext(filename)[-1]))
        cv2.imwrite(output_path, image)

    print("save drawing at {}".format(output_dir))
        
        
        
        

    




